import requests
from requests_toolbelt import MultipartEncoder
import os
import contextlib
import mimetypes
import base64
from dotenv import load_dotenv

load_dotenv()

API_URL = "https://ai.gitee.com/v1/images/edits"
api_key = os.getenv("api_key")
headers = {
	"Authorization": f"Bearer {api_key}" # 使用 Bearer 令牌进行身份验证
}

def query(payload):
	fields = [
		("prompt", payload["prompt"]),
		("model", payload["model"]),
		("steps", str(payload["steps"])),
		("instruction_following_strength", str(payload["instruction_following_strength"])),
		("image_preservation_strength", str(payload["image_preservation_strength"])),
		("refinement_strength", str(payload["refinement_strength"])),
	]
	with contextlib.ExitStack() as stack:
		filepath = payload["image"]
		name = os.path.basename(filepath)
		if filepath.startswith(("http://", "https://")):
			response = requests.get(filepath, timeout=10)
			response.raise_for_status()
			fields.append(("image", (name, response.content, response.headers.get("Content-Type", "application/octet-stream"))))
		else:
			mime_type, _ = mimetypes.guess_type(filepath)
			fields.append(("image", (name, stack.enter_context(open(filepath, "rb")), mime_type or "application/octet-stream")))
		encoder = MultipartEncoder(fields)
		headers["Content-Type"] = encoder.content_type
		response = requests.post(API_URL, headers=headers, data=encoder)
		return response.json()

output = query({
	"image": "/path/to/image.png", # 替换为实际的图像路径
	"prompt": "A sunlit indoor lounge area with a pool containing a flamingo", # 替换为实际的提示文本
	"model": "HiDream-E1-Full",
	"steps": 15,
	"instruction_following_strength": 5,
	"image_preservation_strength": 3,
	"refinement_strength": 0.3
})

for i, image_data in enumerate(output["data"]):
	if "url" in image_data:
		# 从网址下载
		ext = image_data["url"].split('.')[-1].split('?')[0] or "jpg"
		filename = f"HiDream-E1-Full-output-{i}.{ext}"
		response = requests.get(image_data["url"], timeout=30)
		response.raise_for_status()
		with open(filename, "wb") as f:
			f.write(response.content)
		print(f"Downloaded image to {filename}")
	elif "b64_json" in image_data:
		# 解码 Base64
		image_bytes = base64.b64decode(image_data["b64_json"])
		filename = f"HiDream-E1-Full-output-{i}.jpg"
		with open(filename, "wb") as f:
			f.write(image_bytes)
		print(f"Saved image to {filename}")
